# -*- coding:utf-8 -*-

import sys
import os

sys.path.append(os.getcwd())
# sys.path.append("./src")
import base64
import torch
import argparse
import cv2
import numpy as np
from models.factory import build_net
from pybaseutils import image_utils, file_utils


def get_parser():
    image_dir = './data/test_image'
    parser = argparse.ArgumentParser(description='dsfd demo')
    parser.add_argument('--net_type', default='vgg', type=str, choices=['vgg', 'resnet50', 'resnet101', 'resnet152'],
                        help='model for training')
    parser.add_argument('--image_dir', type=str, default=image_dir, help='Directory')
    parser.add_argument('--save_root', type=str, default='out/', help='Directory for detect result')
    parser.add_argument('--model_path', type=str, default='data/pretrained/dsfd_vgg_0.880.pth', help='trained model')
    parser.add_argument('--prob_threshold', default=0.5, type=float,
                        help='Final confidence threshold')
    args = parser.parse_args()
    return args


class_names = ["BACKGROUND", "face"]


class Detector(object):
    def __init__(self, model_path, net_type, prob_threshold=0.6, iou_threshold=0.1, device="cuda:0"):
        self.device = device
        self.class_names = class_names
        self.num_classes = len(self.class_names)
        self.net_type = net_type
        self.model_path = model_path
        self.tok = 750
        self.iou_threshold = iou_threshold
        self.prob_threshold = prob_threshold
        self.net = self.build_net(self.model_path, self.net_type)
        self.net.eval()
        self.net.to(self.device)

    def build_net(self, model_path, net_type):
        """
        :param model_path:  path to model(*.pth) file
        :param net_type:  "RFB" (higher precision) or "slim" (faster)'
        :return:
        """
        net = build_net('test', self.num_classes, net_type)
        net.load_state_dict(torch.load(model_path,map_location="cpu"))
        return net

    def pre_process(self, image, img_mean=[104., 117, 123.], std=None):
        """
        :param image:
        :param img_mean:
        :param std:
        :return:
        """
        height, width, _ = image.shape
        if height > 1500 or width > 1500:
            image = image_utils.resize_image(image, size=(1000,None))
        elif height < 200 or width < 200:
            image = image_utils.resize_image(image, size=(200,None))
        # HWC to CHW
        if len(image.shape) == 3:
            image = np.swapaxes(image, 1, 2)
            image = np.swapaxes(image, 1, 0)
        img_mean = np.array(img_mean, dtype=np.float32)[:, np.newaxis, np.newaxis]
        # RBG to BGR
        image = image[[2, 1, 0], :, :]
        image = image.astype('float32')
        image -= img_mean
        image = image[[2, 1, 0], :, :]
        # image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        image_tensor = torch.from_numpy(image).unsqueeze(0)
        image_tensor = image_tensor.to(self.device)
        return image_tensor

    def forward(self, image_tensor):
        """
        :param image_tensor:
        :return: scores: shape=([1, num_bboxes, num_class])
                 boxes:  shape=([1, num_bboxes, 4]),boxes=[[xmin,ymin,xmax,ymax]]
        """
        with torch.no_grad():
            detections = self.net(image_tensor)
            # detections = y.data
        return detections

    def predict(self, image, prob_threshold, iou_threshold):
        """
        :param rgb_image: RGB Image
        :param top_k: keep top_k results. If k <= 0, keep all the results.
        :param prob_threshold:
        :param iou_threshold:
        :return:
        """
        height, width, _ = image.shape
        image_tensor = self.pre_process(image)
        detections = self.forward(image_tensor)
        bboxes, labels, scores = self.post_process(detections, width, height, prob_threshold, iou_threshold)
        return bboxes, labels, scores

    def detect_image(self, rgb_image, vis=False):
        """
        :param rgb_image:  input RGB Image
        :param vis:
        :return:
        """
        bboxes, labels, scores = self.predict(rgb_image,
                                              iou_threshold=self.iou_threshold,
                                              prob_threshold=self.prob_threshold)

        if vis:
            print("boxes:{}\nlabels:{}\nprobs:{}".format(bboxes, labels, scores))
            self.show_image(rgb_image, bboxes, labels, scores)
        return bboxes, labels, scores

    def post_process(self, detections, width, height, prob_threshold, iou_threshold):
        if not isinstance(detections, np.ndarray):
            detections = detections.cpu().detach().numpy()
        scale = [width, height, width, height]
        scores = []
        bboxes = []
        labels = []
        for i in range(detections.shape[1]):
            j = 0
            while j < self.tok and detections[0, i, j, 0] >= prob_threshold:
                score = detections[0, i, j, 0]
                bboxe = detections[0, i, j, 1:] * scale
                bboxes.append(bboxe)
                scores.append(score)
                labels.append(i)
                j += 1
        bboxes = np.asarray(bboxes)
        bboxes = np.clip(bboxes, 0, max(width, height))
        scores = np.asarray(scores)
        labels = np.asarray(labels)
        return bboxes, labels, scores

    def detect_image_dir(self, image_dir, isshow=True):
        """
        :param image_dir: directory or image file path
        :param isshow:<bool>
        :return:
        """
        if os.path.isdir(image_dir):
            image_list = file_utils.get_files(image_dir, postfix=["*.jpg", "*.png"])
        elif os.path.isfile(image_dir):
            image_list = [image_dir]
        else:
            raise Exception("Error:{}".format(image_dir))
        for img_path in image_list:
            orig_image = cv2.imread(img_path)
            rgb_image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
            # rgb_image = np.array(Image.open(img_path))
            boxes, labels, probs = self.detect_image(rgb_image, vis=isshow)

    def decode_labels(self, labels):
        boxes_name = [self.class_names[i] for i in labels.tolist()]
        return boxes_name

    def show_image(self, image, bboxes: np.asarray, labels, scores):
        """
        :param image: RGB image
        :param bboxes:<np.ndarray>: (num_bboxes, 4), box=[xmin,ymin,xmax,ymax]
        :param scores:<np.ndarray>: (num_bboxes,)
        :param labels:<np.ndarray>: (num_bboxes,)
        :return:
        """
        # boxes_name = [self.class_names[i] for i in labels.tolist()]
        # boxes_name = ["{}:{:3.4f}".format(l, s) for l, s in zip(labels, scores)]
        boxes_name = ["{:3.3f}".format(s) for s in scores]
        image = file_utils.draw_image_bboxes_text(image, bboxes, boxes_name, color=(255, 0, 0))
        # image = image_processing.draw_image_boxes(image, bboxes)
        bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        # self.save_image(bgr_image)
        # cv2.imwrite("result.jpg", bgr_image)
        cv2.imshow("Det", bgr_image)
        # cv2.imshow("Det", image_processing.resize_image(bgr_image, 1080))
        cv2.waitKey(0)


if __name__ == '__main__':
    args = get_parser()
    print(args)
    det = Detector(model_path=args.model_path,
                   net_type=args.net_type,
                   prob_threshold=args.prob_threshold)
    det.detect_image_dir(image_dir=args.image_dir)
